/**
 * 
 */
package com.gragra.sampling.nonproj;
import java.util.concurrent.atomic.AtomicInteger;
import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.ints.IntIterator;
import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
import it.unimi.dsi.fastutil.ints.IntSet;
import com.gragra.sampling.CRPProcess;
import com.gragra.sampling.LogarithmicFactorialPlusAlphaCash;
import com.gragra.sampling.PopulationSamplingStructure;
import com.gragra.sampling.binary.HeadAssignmentStructure;
import com.gragra.sampling.nonproj.structures.nonprojective.NonProjDepdTree;
import com.gragra.util.OpenUnitElementCounter;
/**
 * @author Christoph Teichmann
 * 30.09.2013
 */
public class DistHeadModel extends CRPProcess implements ProposalDependencyAssigner
{
	/**
	 * 
	 */
	private static final int	DISTANCES	= 10;
	/**
	 * 
	 */
	private static final double	INCREASE	= 1.0;
	/**
	 * 
	 */
	private static final double	REDUCTION	= -1.0;
	/**
	 * 
	 */
	private static final int ROOTNODETAG	= -1;
	/**
	 * 
	 */
	private static final int MAXCHILDREN	= 10;
	/**
	 * 
	 */
	private final OpenUnitElementCounter[] holder;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash distUnit;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash distElement;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash tagUnit;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash tagElement;
	/**
	 * 
	 */
	private int free = 0;
	/**
	 * 
	 * @return
	 */
	protected int getFreeUnit()
	{return free++;}
	/**
	 * 
	 */
	private final IntSet tUnits ;
	/**
	 * 
	 */
	private final IntSet dUnits;
	/**
	 * 
	 */
	private final Int2IntOpenHashMap tagTagUnits = new Int2IntOpenHashMap();
	/**
	 * 
	 */
	private final Int2IntOpenHashMap tagDistUnits = new Int2IntOpenHashMap();
	/**
	 *
	 */
	private final Int2IntOpenHashMap fertUnits = new Int2IntOpenHashMap();	
	/**
	 * @param population
	 */
	public DistHeadModel(int population, IntSet tags, double tagTagAlpha, double tagDistAlpha,
									double fertAlpha)
	{
		super(population);
		IntIterator iit = tags.iterator();
		while(iit.hasNext())
		{tagTagUnits.put(iit.nextInt(), getFreeUnit());}
		tagTagUnits.put(ROOTNODETAG, getFreeUnit());
		tUnits = new IntOpenHashSet(tagTagUnits.values());
		iit = tags.iterator();
		while(iit.hasNext())
		{
			int tag = iit.nextInt();
			tagDistUnits.put(tag, getFreeUnit());
			fertUnits.put(tag, getFreeUnit());
		}
		fertUnits.put(ROOTNODETAG, getFreeUnit());
		fUnits = new IntOpenHashSet(fertUnits.values());
		dUnits = new IntOpenHashSet(tagDistUnits.values());
		this.distElement = new LogarithmicFactorialPlusAlphaCash(tagDistAlpha, population);
		this.distUnit = new LogarithmicFactorialPlusAlphaCash(tagDistAlpha*((DISTANCES*2)+1), population);
		this.tagElement = new LogarithmicFactorialPlusAlphaCash(tagTagAlpha, population);
		this.tagUnit = new LogarithmicFactorialPlusAlphaCash(tagTagAlpha*tags.size(), population);
		this.fertUnit = new LogarithmicFactorialPlusAlphaCash(fertAlpha*((MAXCHILDREN+1)), population);
		this.fertElement = new LogarithmicFactorialPlusAlphaCash(fertAlpha, population);
		this.holder = new OpenUnitElementCounter[population];
		this.child = new IntArrayList[population];
		for(int i=0;i<population;++i)
		{
			holder[i] = new OpenUnitElementCounter();
			child[i] = new IntArrayList();
		}
	}
	/**
	 * 
	 */
	private final AtomicInteger at = new AtomicInteger();
	/**
	 * 
	 */
	private final IntSet fUnits;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash	fertElement;
	/**
	 * 
	 */
	private LogarithmicFactorialPlusAlphaCash	fertUnit;
	/* (non-Javadoc)
	 * @see com.gragra.sampling.PopulationProbabilityAssigner#makeLogProb(com.gragra.sampling.PopulationSamplingStructure, int)
	 */
	@Override
	public double makeLogProb(PopulationSamplingStructure pss, int thread)
	{
		HeadAssignmentStructure has = (HeadAssignmentStructure) pss;
		this.at.addAndGet(has.getLength());
		double d = 0.0;
		OpenUnitElementCounter oue = getSeen(thread);
		IntArrayList dal = this.makeChildren(has, thread);
		for(int i=0;i<has.getLength();++i)
		{
			addFeatures(i,oue,has,thread,1.0,dal);
			d += this.makeDistanceAndOtherConstraints(i,has,thread);
		}
		d += this.makeLogProb(oue,thread);
		return d;
	}
	/**
	 * @param thread
	 * @return
	 */
	protected OpenUnitElementCounter getSeen(int thread)
	{
		OpenUnitElementCounter oue = this.holder[thread];
		oue.clear();
		return oue;
	}
	/**
	 * 
	 * @param i
	 * @param oue
	 * @param has
	 * @param thread
	 * @param amount 
	 */
	protected void addFeatures(int pos, OpenUnitElementCounter oue,
			HeadAssignmentStructure has, int thread, double amount, IntArrayList children)
	{
		int head = has.getHeadPosition(pos, thread);
		int headTag = ROOTNODETAG;
		int dist = 0;
		if(head >= 0 && head != has.getDependencyRootPosition())
		{
			headTag = has.getTag(head);
			dist = Math.min(Math.max(head-pos,-DISTANCES),DISTANCES);
		}
		int tag = has.getTag(pos);
		int hhUnit = getHeadToHeadUnit(tag);
		int dUnit = getDistUnit(tag);
		int fUnit = getFertilityUnit(tag);
		oue.add(fUnit, makeFertility(pos, children),amount);
		if(pos == 0)
		{
			fUnit = getFertilityUnit(ROOTNODETAG);
			oue.add(fUnit,makeFertility(has.getDependencyRootPosition(), children),amount);
		}
		if(head >= 0)
		{
			oue.add(hhUnit, headTag, amount);
			oue.add(dUnit, dist, amount);
		}
	}
	/**
	 * @param pos
	 * @param children
	 * @return
	 */
	protected int makeFertility(int pos, IntArrayList children)
	{
		int ch = pos < children.size() ? children.getInt(pos) : 0;
		ch = Math.min(Math.max(ch, -MAXCHILDREN), MAXCHILDREN);
		return ch;
	}
	/**
	 * 
	 * @param tag
	 * @return
	 */
	private int getFertilityUnit(int tag)
	{return this.fertUnits.get(tag);}
	/**
	 * @param tag
	 * @return
	 */
	protected int getHeadToHeadUnit(int tag)
	{return this.tagTagUnits.get(tag);}
	/**
	 * @param tag
	 * @return
	 */
	protected int getDistUnit(int tag)
	{return this.tagDistUnits.get(tag);}
	/**
	 * 
	 * @param i
	 * @param has
	 * @param thread
	 * @return
	 */
	protected double makeDistanceAndOtherConstraints(int pos, HeadAssignmentStructure has,
			int thread)
	{return 0.0;}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.PopulationProbabilityAssigner#add(com.gragra.sampling.PopulationSamplingStructure, int)
	 */
	@Override
	public void add(PopulationSamplingStructure pss, int thread)
	{justHandleFeatures(pss, thread, INCREASE);}
	/**
	 * @param pss
	 * @param thread
	 * @param dir
	 */
	protected void justHandleFeatures(PopulationSamplingStructure pss,
			int thread, double dir)
	{
		HeadAssignmentStructure has = (HeadAssignmentStructure) pss;
		OpenUnitElementCounter oue = super.getGiven(thread);
		IntArrayList ial = makeChildren(has,thread);
		for(int i=0;i<has.getLength();++i)
		{addFeatures(i,oue,has,thread,dir, ial);}
	}
	/**
	 * 
	 * @param has
	 * @param thread
	 * @return
	 */
	private IntArrayList makeChildren(HeadAssignmentStructure has, int thread)
	{
		IntArrayList list = child[thread];
		list.clear();
		for(int i=0;i<has.getLength();++i)
		{
			int head = has.getHeadPosition(i, thread);
			if(head == -1)
			{continue;}
			while(list.size() <= head)
			{list.add(0);}
			list.set(head, list.getInt(head)+1);
		}
		return list;
	}
	/**
	 * 
	 */
	private final IntArrayList[] child;
	/* (non-Javadoc)
	 * @see com.gragra.sampling.PopulationProbabilityAssigner#remove(com.gragra.sampling.PopulationSamplingStructure, int)
	 */
	@Override
	public void remove(PopulationSamplingStructure pss, int thread)
	{this.justHandleFeatures(pss, thread, REDUCTION);}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.PopulationProbabilityAssigner#getLocalLogLike(com.gragra.sampling.PopulationSamplingStructure, int)
	 */
	@Override
	public double getLocalLogLike(PopulationSamplingStructure pss, int thread)
	{
		HeadAssignmentStructure has = (HeadAssignmentStructure) pss;
		double d = 0.0;
		for(int i=0;i<has.getLength();++i)
		{d += this.makeDistanceAndOtherConstraints(i,has,thread);}
		return d;
	}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.CRPProcess#getUnitCash(int)
	 */
	@Override
	protected LogarithmicFactorialPlusAlphaCash getUnitCash(int unit)
	{
		if(this.dUnits.contains(unit))
		{return this.distUnit;}
		else if (this.tUnits.contains(unit)) 
		{return this.tagUnit;}
		else if(this.fUnits.contains(unit))
		{return this.fertUnit;}
		return null;
	}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.CRPProcess#getElementCash(int, int)
	 */
	@Override
	protected LogarithmicFactorialPlusAlphaCash getElementCash(int unit,
			int element)
	{
		if(this.dUnits.contains(unit))
		{return this.distElement;}
		else if (this.tUnits.contains(unit)) 
		{return this.tagElement;}
		else if(this.fUnits.contains(unit))
		{return this.fertElement;}
		return null;
	}
	@Override
	public double getNumberOfEvaluations()
	{return this.at.get();}
	@Override
	public double makeProposalLogProb(int pos, int head,
			NonProjDepdTree tree, int thread)
	{
		this.auxs.addAndGet(tree.getLength());
		double d = 0.0;
		OpenUnitElementCounter oue = getSeen(thread);
		IntArrayList dal = this.makeChildren(tree, thread);
		for(int i=0;i<tree.getLength();++i)
		{
			addFeatures(i,oue,tree,thread,1.0,dal);
			d += this.makeDistanceAndOtherConstraints(i,tree,thread);
		}
		d += this.makeLogProb(oue,thread);
		return d;
	}
	@Override
	public double makeChartProposalProb(
			NonProjDepdTree tree,
			int headCandidate, int child, int thread)
	{
		this.auxs.incrementAndGet();
		OpenUnitElementCounter oue = this.getSeen(thread);
		int tag = tree.getTag(child);
		int headTag = headCandidate == tree.getRootPosition() ? ROOTNODETAG : tree.getTag(headCandidate);
		int hunit = this.getHeadToHeadUnit(tag);
		oue.add(hunit, headTag, 1.0);
		int dist = headCandidate == tree.getRootPosition() ? 0 : Math.min(Math.max(headCandidate-child,-DISTANCES),DISTANCES);
		int dUnit = this.getDistUnit(tag);
		oue.add(dUnit, dist, 1.0);
		return this.makeLogProb(oue, thread);
	}
	@Override
	public double getNumberOfAuxEvals()
	{return auxs.get();}
	/**
	 * 
	 */
	private final AtomicInteger auxs = new AtomicInteger();
	@Override
	public void reduceStructuralAnnealing(int p)
	{}
}